#  -*- coding: utf-8 -*-

import traceback
import matplotlib.pyplot as plt
import numpy as np
from pymongo import UpdateOne
import pandas as pd
from data.data_module import DataModule
from trading.signal.computer.base_signal_computer import BaseSignalComputer
from util.stock_util import get_all_codes
import mpl_finance as mpf

class BollSignalComputer(BaseSignalComputer):
    def __init__(self):
        BaseSignalComputer.__init__(self, 'boll_signal')

    def compute(self, begin_date, end_date):
        """
        计算指定日期内的信号
        :param begin_date: 开始日期
        :param end_date: 结束日期
        """
        # all_codes = get_all_codes('2018-08-24')
        all_codes=['601398']
        dm = DataModule()

        N = 20
        k = 2
        mistake=0.1
        for code in all_codes[:]:
            print (code)
            try:
                df_daily = dm.get_k_data(code, autype='kba', begin_date=begin_date, end_date=end_date,trade=True)
                # 计算MB，盘后计算，这里用当日的Close
                df_daily['MID'] = df_daily['close'].rolling(N).mean().astype('float')
                # 计算STD20
                df_daily['std'] = df_daily['close'].rolling(N).std().astype('float')
                # 计算UP
                df_daily['UP'] = df_daily['MID'] + k * df_daily['std']
                # 计算down
                df_daily['DOWN'] = df_daily['MID'] - k * df_daily['std']
                df_daily['close']=df_daily['close'].astype('float')
                df_daily['date']=df_daily.index
                a = pd.DataFrame(df_daily, columns=['date','UP','MID','DOWN','close'])
                a.dropna(inplace=True)
                b= pd.DataFrame(df_daily, columns=['open', 'high', 'low', 'close'])
                b['open']=b['open'].astype('float64')
                b['high'] = b['high'].astype('float64')
                b['low'] = b['low'].astype('float64')
                b['close'] = b['close'].astype('float64')
                candleData = np.column_stack([list(range(len(df_daily['date']))), b])
                fig = plt.figure(figsize=(10, 6))
                ax = fig.add_axes([0.1, 0.3, 0.8, 0.6])
                mpf.candlestick_ohlc(ax, candleData, width=0.5, colorup='r', colordown='b')
                plt.plot(df_daily['DOWN'])
                plt.plot(df_daily['UP'])
                plt.plot(df_daily['MID'])
                plt.xticks([])
                plt.show()
            except:
                traceback.print_exc()
            # break

if __name__ == '__main__':
    BollSignalComputer().compute(begin_date='2018-01-01', end_date='2018-08-30')
